import numpy as np

from pytorch_grad_cam.base_cam import BaseCAM
from typing import Optional, Callable


class GradCAM(BaseCAM):
    def __init__(
        self, model, target_layers, reshape_transform: Optional[Callable] = None
    ):
        super(GradCAM, self).__init__(model, target_layers, reshape_transform)

    def get_cam_weights(
        self, input_tensor, target_layer, target_category, activations, grads
    ):
        """
        Returns:
            weights: [B, 768] for 2D image, the weight for different channels
        """
        # 2D image
        if len(grads.shape) == 4:
            return np.mean(grads, axis=(2, 3))

        # 3D image
        elif len(grads.shape) == 5:
            return np.mean(grads, axis=(2, 3, 4))

        else:
            raise ValueError(
                "Invalid grads shape."
                "Shape of grads should be 4 (2D image) or 5 (3D image)."
            )
